
package com.studiofortress.sf.util.collections;

import java.util.ArrayList;
import java.util.Collection;
import java.util.ConcurrentModificationException;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * This is an implementation of a HashMap based on java.util.HashMap, however
 * it is designed for times when you very often insert and remove values.
 *
 * Internally old data structures are cached when values are removed and reused
 * when values are needed. This is slightly faster, but more importantly can
 * give constant memory usage and reduce garbage collection.
 *
 * To clear the internal cache you should use the clearCache method. Otherwise
 * it will never be cleared (even if you call clear). Values and keys are not
 * cached, when they are removed they are fully removed.
 *
 * @author Joseph Lenton - JosephLenton@StudioFortress.com
 */
public class CachingHashMap<A, B> implements Map<A, B>, CallbackIterable<B>
{
    private static final int MAXIMUM_CAPACITY      = 1 << 30;
    private static final int DEFAULT_CAPACITY      = 16;
    private static final float DEFAULT_LOAD_FACTOR = 0.75f;

    /**
     * @param key The object to get the hash value for, cannot be null.
     * @return A hash code generated from the given object.
     */
    private final static int hash(final Object key)
    {
        return hash( key.hashCode() );
    }

    /**
     * The key to this hash function is that it takes the hash value given and
     * turns into one that is more spread out.
     * @param index The value to hash.
     * @return The given value hashed so it's more spread out across the possible values.
     */
    private final static int hash(int value)
    {
        value ^= (value >>> 20) ^ (value >>> 12);
        return value ^ (value >>> 7) ^ (value >>> 4);
    }

    private final static int indexOf(final int hash, final int length)
    {
        return hash & (length-1);
    }

    private int size;
    private float loadFactor;
    private int threshold;
    private CachingEntry<A, B>[] entries;
    private final ObjectStack< CachingEntry<A, B> > entriesCache;

    // Every time an operation occurres this is incremented.
    // A copy is given to new iterators so they can check if operations have
    // occurred after being created.
    private int failFast;
    
    public CachingHashMap()
    {
        this( DEFAULT_CAPACITY, DEFAULT_LOAD_FACTOR );
    }
    
    public CachingHashMap(int initialCapacity)
    {
        this( initialCapacity, DEFAULT_LOAD_FACTOR );
    }
    
    public CachingHashMap(int initialCapacity, float loadFactor)
    {
        if (initialCapacity < 0) {
            throw new IllegalArgumentException("Illegal initial capacity: " + initialCapacity);
        } else  if (initialCapacity > MAXIMUM_CAPACITY) {
            initialCapacity = MAXIMUM_CAPACITY;
        } else if (loadFactor <= 0 || Float.isNaN(loadFactor)) {
            throw new IllegalArgumentException("Illegal load factor: " + loadFactor);
        }
        
        initialCapacity = getPowerOfTwoCapacity( initialCapacity );
        
        this.size = 0;
        this.loadFactor = loadFactor;
        entries = ( CachingEntry<A, B>[] ) new CachingEntry[initialCapacity];
        threshold = (int) ( initialCapacity*loadFactor );
        entriesCache  = new ObjectStack< CachingEntry<A, B> >();

        failFast = 0;
    }

    /**
     * Iterates over each of the values stored in this map.
     * @param callback The callback to run each element in this map, cannot be null.
     */
    @Override
    public void iterate(CallbackIterator<B> callback)
    {
        if ( callback == null ) {
            throw new IllegalArgumentException("The given iterator cannot be null.");
        }

        final CachingEntry<A, B>[] table = entries;

        for (int i = 0; i < table.length; i++) {
            for (CachingEntry<A, B> e = table[i]; e != null; e = e.next) {
                callback.iterate( e.value );
            }
        }
    }

    /**
     * Iterates over each of the values stored in this map. This one has the
     * option of only partially iterating.
     * @param callback The callback to run each element in this map, cannot be null.
     */
    @Override
    public void iteratePartial(CallbackPartialIterator<B> callback)
    {
        if ( callback == null ) {
            throw new IllegalArgumentException("The given iterator cannot be null.");
        }

        final CachingEntry<A, B>[] table = entries;
        callback.startIteration();

        for (int i = 0; i < table.length; i++) {
            for (CachingEntry<A, B> e = table[i]; e != null; e = e.next) {
                if ( callback.isIterating() ) {
                    callback.iterate( e.value );
                } else {
                    return;
                }
            }
        }
    }

    /**
     * Iterates over each of the keys stored in this map.
     * @param callback The callback to run each key in this map against, cannot be null.
     */
    public void iterateKeys(CallbackIterator<A> callback)
    {
        if ( callback == null ) {
            throw new IllegalArgumentException("The given iterator cannot be null.");
        }

        final CachingEntry<A, B>[] table = entries;

        for (int i = 0; i < table.length; i++) {
            for (CachingEntry<A, B> e = table[i]; e != null; e = e.next) {
                callback.iterate( e.key );
            }
        }
    }

    /**
     * Iterates over each of the keys stored in this map.
     * @param callback The callback to run each key in this map against, cannot be null.
     */
    public void iterateKeysPartial(CallbackPartialIterator<A> callback)
    {
        if ( callback == null ) {
            throw new IllegalArgumentException("The given iterator cannot be null.");
        }

        final CachingEntry<A, B>[] table = entries;
        callback.startIteration();

        for (int i = 0; i < table.length; i++) {
            for (CachingEntry<A, B> e = table[i]; e != null; e = e.next) {
                if ( callback.isIterating() ) {
                    callback.iterate( e.key );
                } else {
                    return;
                }
            }
        }
    }
    
    private int getPowerOfTwoCapacity(int initialCapacity)
    {
        int capacity = 1;
        while (capacity < initialCapacity) {
            capacity <<= 1;
        }
        return capacity;
    }

    public B put(A key, B value)
    {
        ++failFast;

        final int hash  = hash( key );
        final int index = indexOf( hash, entries.length );

        for ( CachingEntry<A, B> entry = entries[ index ];
                entry != null;
                entry = entry.next)
        {
            if (
                    entry.hash == hash &&
                    (entry.key == key || entry.key.equals(key)) )
            {
                B oldValue = entry.value;
                entry.value = value;
                return oldValue;
            }
        }
        
        addEntry( hash, key, value, index );
        return null;
    }

    private void addEntry(int hash, A key, B value, int index)
    {
	final CachingEntry<A, B> e = entries[ index ];
        entries[ index ] = newEntry( hash, key, value, e );
        
        if ( size++ >= threshold ) {
            resize( 2 * entries.length );
        }
    }

    private void resize(int newSize)
    {
        if ( entries.length == MAXIMUM_CAPACITY ) {
            threshold = Integer.MAX_VALUE;
            return;
        }
        
        final CachingEntry[] newTable = new CachingEntry[ newSize ];
        transfer( newTable );
        entries = newTable;
        threshold = (int) ( newSize * loadFactor );
    }

    private void transfer(CachingEntry<A, B>[] newTable)
    {
        final CachingEntry<A, B>[] src = entries;
        final int newCapacity = newTable.length;
        
        for (int j = 0; j < src.length; j++) {
            CachingEntry<A, B> e = src[j];
            
            if (e != null) {
                src[j] = null;
                do {
                    final CachingEntry<A, B> next = e.next;
                    final int i = indexOf( e.hash, newCapacity );
                    
                    e.next = newTable[i];
                    newTable[i] = e;
                    e = next;
                } while (e != null);
            }
        }
    }
    
    @Override
    public B get(final Object key)
    {
        final int hash  = hash( key );
        
        for ( CachingEntry<A, B> entry = entries[ indexOf(hash, entries.length) ];
                entry != null;
                entry = entry.next)
        {
            if (
                    entry.hash == hash &&
                    (entry.key == key || entry.key.equals(key)) )
            {
                return entry.value;
            }
        }
        
        return null;
    }

    @Override
    public B remove(Object key)
    {
        ++failFast;

        int hash = hash( key );
        int i = indexOf(hash, entries.length);
        
        CachingEntry<A, B> prev = entries[i];
        CachingEntry<A, B> entry = prev;

        while (entry != null) {
            CachingEntry<A, B> next = entry.next;

            if (entry.hash == hash && entry.key == key) {
                size--;

                if (prev == entry) {
                    entries[i] = next;
                } else {
                    prev.next = next;
                }

                B value = entry.value;
                cacheEntry( entry );
                return value;
            }

            prev = entry;
            entry = next;
        }

        return null;
    }
    
    public boolean isEmpty()
    {
        return size() == 0;
    }
    
    public int size()
    {
        return size;
    }

    public void clear()
    {
        ++failFast;
        
        final CachingEntry<A, B>[] table = entries;
        for (int i = 0; i < table.length; i++) {
            CachingEntry<A, B> next;

            for (CachingEntry<A, B> e = table[i]; e != null; e = next) {
                next = e.next;
                cacheEntry( e );
            }
            
            table[i] = null;
        }
        
        size = 0;
    }

    /**
     * Clears the internal cache of data structures that this hash map uses.
     * This will not clear this map.
     */
    public void clearCache()
    {
        entriesCache.clear();
    }

    /**
     * @return The number of items currently cached in the map.
     */
    public int cacheSize()
    {
        return entriesCache.size();
    }

    /**
     * This will return a new entry from the cache of entries. The entry will be
     * initialised using the values given. If an entry was not found, then a new
     * entry will be created.
     * @param <A>
     * @param <B>
     * @param hash
     * @param key
     * @param item
     * @param next
     * @return
     */
    private final <A, B> CachingEntry<A, B> newEntry(
            final int hash, final A key, final B item, final CachingEntry<A, B> next)
    {
        if ( entriesCache.size() == 0 ) {
            return new CachingEntry<A, B>( hash, key, item, next );
        } else {
            return ( (CachingEntry<A, B>) entriesCache.pop() ).initialize( hash, key, item, next );
        }
    }
    
    /**
     * Stores the given node in the cache. After this is called the entry should
     * never be used again, unless it is returned from newEntry.
     * It's references to it's value and the entry will be set to null for you.
     * @param node The node to store, cannot be null.
     */
    private void cacheEntry(CachingEntry<A, B> node)
    {
        node.value = null;
        node.next = null;
        entriesCache.push( node );
    }

    /**
     * Same as clear only it performs no caching of the internal data
     * structures. As a result it caches a lot faster then the other clear
     * method.
     */
    public void clearFast()
    {
        ++failFast;

        for (int i = 0; i < entries.length; i++) {
            entries[i] = null;
        }
        
        size = 0;
    }

    /**
     * The keys returned are in no particular order and are not guaranteed to
     * always be returned in the same order.
     * @return A set containing all of the keys stored in this map.
     */
    public Set<A> keySet()
    {
        Set<A> keys = new HashSet<A>();

        for (int i = 0; i < entries.length; i++) {
            for (CachingEntry<A, B> entry = entries[i]; entry != null; entry = entry.next) {
                keys.add( entry.key );
            }
        }
        
        return keys;
    }
    
    public boolean containsKey(final Object key)
    {
        final int hash  = hash( key );
        final int index = indexOf( hash, entries.length );

        for ( CachingEntry<A, B> entry = entries[ index ];
                entry != null;
                entry = entry.next)
        {
            if (entry.hash == hash && entry.key == key) {
                return true;
            }
        }
        
        return false;
    }

    public boolean containsValue(Object value)
    {
        if (value == null) {
            for (int i = 0; i < entries.length; i++) {
                for (CachingEntry<A, B> entry = entries[i]; entry != null; entry = entry.next) {
                    if (entry.value == null) {
                        return true;
                    }
                }
            }
        } else {
            for (int i = 0; i < entries.length; i++) {
                for (CachingEntry<A, B> entry = entries[i]; entry != null; entry = entry.next) {
                    if (entry.value.equals(value)) {
                        return true;
                    }
                }
            }
        }
        
        return false;
    }

    /**
     * Puts all the entries from the given map into this map.
     * @param map The map of entries to store the values from, cannot be null.
     */
    public void putAll(Map<? extends A, ? extends B> map)
    {
        if (map == null) {
            throw new IllegalArgumentException("The given map cannot be null");
        }
        
        ++failFast;

        for (Map.Entry<? extends A, ? extends B> entry : map.entrySet()) {
            put( entry.getKey(), entry.getValue() );
        }
    }

    /**
     * The order of the collection returned is in no particular order, and it
     * should not be presumed that it will always return values in the same
     * order.
     * @return A collection containing all the values stored in this map.
     */
    public Collection<B> values()
    {
        final List<B> values = new ArrayList<B>( size() );

        for (int i = 0; i < entries.length; i++) {
            for (CachingEntry<A, B> entry = entries[i]; entry != null; entry = entry.next) {
                values.add( entry.value );
            }
        }

        return values;
    }

    public Set<Map.Entry<A, B>> entrySet()
    {
        final Set<Map.Entry<A, B>> entrySets = new HashSet<Map.Entry<A, B>>( size() );
        
        for (int i = 0; i < entries.length; i++) {
            for (CachingEntry<A, B> entry = entries[i]; entry != null; entry = entry.next) {
                entrySets.add( new MapEntry<A, B>(entry) );
            }
        }
        
        return entrySets;
    }
    
    public Iterator<A> getKeyIterator()
    {
        return new KeyIterator( failFast );
    }
    
    private static final class CachingEntry<A, B>
    {
        private int      hash;
        private A        key;
        private B        value;
        private CachingEntry<A, B> next;
        
        public CachingEntry(int hash, A key, B value, CachingEntry<A, B> next)
        {
            initialize( hash, key, value, next );
        }

        public CachingEntry<A, B> initialize(int hash, A key, B value, CachingEntry<A, B> next)
        {
            this.hash  = hash;
            this.key   = key;
            this.value = value;
            this.next  = next;
            
            return this;
        }
    }
    
    private class MapEntry<A, B> implements Map.Entry<A, B>
     {
        private final CachingEntry<A, B> entry;
        private final A key;
        private B value;

        public MapEntry(CachingEntry<A, B> entry)
        {
            this.entry = entry;
            this.key = entry.key;
            this.value = entry.value;
        }

        public A getKey()
        {
            return key;
        }

        public B getValue()
        {
            return value;
        }
        
        /**
         * Unsupported.
         * @param value
         * @return
         */
        public B setValue(B value)
        {
            ++failFast;
            
            // check if the old entry has been cached or re-used for something else
            if (entry.key != key) {
                throw new IllegalStateException("The underlying map has changed, this update cannot be made.");
            }

            final B oldValue = entry.value;

            entry.value = value;
            this.value = value;

            return oldValue;
        }
    }

    private class KeyIterator implements Iterator<A>
    {
        private CachingEntry<A, B> lastNode;
        private CachingEntry<A, B> node;
        private int i;

        private final int failState;

        public KeyIterator(int failFastCode)
        {
            this.failState = failFastCode;
            
            while (node == null && i < entries.length) {
                node = entries[i];
                ++i;
            }
        }
        
        private final void checkFail()
        {
            if ( failFast != failState ) {
                throw new ConcurrentModificationException("The list has been altered outside of this iterator.");
            }
        }
        
        public boolean hasNext()
        {
            checkFail();

            return node != null;
        }
        
        public A next()
        {
            checkFail();

            lastNode = node;
            node = node.next;
            while (node == null && i < entries.length) {
                node = entries[i];
                ++i;
            }

            return lastNode.key;
        }
        
        public void remove()
        {
            checkFail();
            
            throw new UnsupportedOperationException("remove is not built yet");
        }
    }
}
